"""
ground-truth files

image_1.txt
    tvmonitor 2 10 173 238
    book 439 157 556 241
    book 437 246 518 351 difficult
    pottedplant 272 190 316 259

detection-result files

image_1.txt
    tvmonitor 0.471781 0 13 174 244
    cup 0.414941 274 226 301 265
    book 0.460851 429 219 528 247
    chair 0.292345 0 199 88 436
    book 0.269833 433 260 506 336

"""
import os
import numpy as np
import cv2
from collections import defaultdict
import matplotlib.pyplot as plt


class MetricsVOC(object):
    def __init__(self, iou_threshold=0.5, no_difficult=True):
        self.iou_threshold = iou_threshold
        self.no_difficult = no_difficult

    def voc_ap(self, gt_folder, det_folder):
        gt_boxes, class_names, num_per_class = self.load_gts(gt_folder)
        det_boxes = self.load_dets(det_folder)

        aps = []
        for c in sorted(class_names):
            dets = det_boxes[c]
            gts = gt_boxes[c]
            num_gts = num_per_class[c]

            # sort the detections in descending order
            dets = sorted(dets, key=lambda x: x[4], reverse=True)
            # print(dets)

            num_dets = len(dets)
            TP = np.zeros(num_dets)
            FP = np.zeros(num_dets)

            for det_idx, det in enumerate(dets):
                fname = det[-1]
                if fname not in gts.keys():
                    # no this category of gt found in this image, 
                    # the detection is false positive.
                    FP[det_idx] = 1
                    continue

                gts_for_file = gts[fname]

                max_iou = 0.0
                max_idx = -1
                for gt_idx, gt in enumerate(gts_for_file):
                    iou = self.box_overlap(det, gt)
                    if iou > max_iou:
                        max_iou = iou
                        max_idx = gt_idx
                
                if max_iou >= self.iou_threshold:
                    if gts_for_file[gt_idx][4] == 1:
                        # this gt is already matched
                        FP[det_idx] = 1
                    else:
                        # good detection
                        TP[det_idx] = 1
                        gts_for_file[gt_idx][4] = 1
                else:
                    # iou less than threshold
                    FP[det_idx] = 1
            
            # calculate
            acc_FP = np.cumsum(FP)
            acc_TP = np.cumsum(TP)
            recall = acc_TP / num_gts
            precision = np.divide(acc_TP, (acc_TP + acc_FP))

            # calculate ap
            ap = self.class_ap(recall, precision)
            print('{:02f}% = {}'.format(ap * 100, c))

            aps.append({'class_name': c, 'ap': ap})
        print('----------------------------------------')
        aps = sorted(aps, key=lambda x: x['ap'], reverse=True)

        all_ap = []
        for item in aps:
            class_name = item['class_name']
            ap = item['ap']
            all_ap.append(ap)
            # print('{}: {:.02f}%'.format(class_name, ap * 100))

        mAP = sum(all_ap) / len(all_ap)
        print('mAP: {:.03f}%'.format(mAP * 100))

        return mAP
    
    def class_ap(self, recall, precision):
        mrec = []
        mrec.append(0)
        for v in recall:
            mrec.append(v)
        mrec.append(1.0)
        mpre = []
        mpre.append(0)
        for v in precision:
            mpre.append(v)
        mpre.append(0)

        for i in range(len(mpre) - 2, -1, -1):
            mpre[i] = max(mpre[i], mpre[i + 1])
        i_list = []
        for i in range(1, len(mrec)):
            if mrec[i] != mrec[i - 1]:
                i_list.append(i)
        
        ap = 0.0
        for i in i_list:
            ap += ((mrec[i] - mrec[i - 1]) * mpre[i])
        return ap

    def box_overlap(self, box1, box2):
        xmin1, ymin1, xmax1, ymax1 = box1[:4]
        xmin2, ymin2, xmax2, ymax2 = box2[:4]
        area1 = (xmax1 - xmin1) * (ymax1 - ymin1)
        area2 = (xmax2 - xmin2) * (ymax2 - ymin2)
        xmax = min(xmax1, xmax2)
        ymax = min(ymax1, ymax2)
        xmin = max(xmin1, xmin2)
        ymin = max(ymin1, ymin2)
        w = max(0, xmax - xmin)
        h = max(0, ymax - ymin)
        inter = w * h
        return inter * 1.0 / (area1 + area2 - inter)
    
    def view(self, image_folder, gt_folder, det_folder, wait_in_milliseconds=0):
        """
        view gt and dets per image
        """
        for image_name in os.listdir(image_folder):
            base_name = '.'.join(image_name.split('.')[:-1])
            gt_path = os.path.join(gt_folder, '{}.txt'.format(base_name))
            det_path = os.path.join(det_folder, '{}.txt'.format(base_name))
            image_path = os.path.join(image_folder, image_name)
            image = cv2.imread(image_path)

            gt_boxes = []
            with open(gt_path, 'r') as f:
                for line in f.readlines():
                    vec = line.strip().split()
                    class_name = vec[0]
                    box = [int(v) for v in vec[1:5]]
                    # gt_boxes.append([box[0], box[1], box[2], box[3], class_name])
                    cv2.rectangle(image, (box[0], box[1]), (box[2], box[3]), (0, 255, 0), 2)
            det_boxes = []
            with open(det_path, 'r') as f:
                for line in f.readlines():
                    vec = line.strip().split()
                    class_name = vec[0]
                    conf = float(vec[1])
                    box = [int(v) for v in vec[2:6]]
                    # det_boxes.append([box[0], box[1], box[2], box[3], conf, class_name])
                    cv2.rectangle(image, (box[0], box[1]), (box[2], box[3]), (0, 0, 255), 2)
            cv2.imshow('display', image)
            cv2.waitKey(wait_in_milliseconds)
            
    def load_gts(self, gt_folder):
        gt_files = [f for f in os.listdir(gt_folder) if f.endswith('.txt')]
        # gt_files.sort()

        class_names = []
        num_per_class = {}
        gt_boxes = {}

        for f in gt_files:
            fname = f.replace('.txt', '')
            fpath = os.path.join(gt_folder, f)
            with open(fpath, 'r') as fp:
                for line in fp.readlines():
                    if 'difficult' in line and self.no_difficult:
                        continue
                    
                    vec = line.strip().split()
                    class_name = vec[0]
                    x1 = int(vec[1])
                    y1 = int(vec[2])
                    x2 = int(vec[3])
                    y2 = int(vec[4])
                    box = [x1, y1, x2, y2, 0]
                    
                    if class_name not in class_names:
                        class_names.append(class_name)
                        gt_boxes[class_name] = defaultdict(list)
                        num_per_class[class_name] = 0
                    gt_boxes[class_name][fname].append(box)
                    num_per_class[class_name] += 1
        
        return gt_boxes, class_names, num_per_class

    def load_dets(self, det_folder):
        det_files = [f for f in os.listdir(det_folder) if f.endswith('.txt')]
        # det_files.sort()

        det_boxes = defaultdict(list)

        for f in det_files:
            fname = f.replace('.txt', '')
            fpath = os.path.join(det_folder, f)
            with open(fpath, 'r') as fp:
                for line in fp.readlines():
                    vec = line.strip().split()
                    class_name = vec[0]
                    conf = float(vec[1])
                    x1 = int(vec[2])
                    y1 = int(vec[3])
                    x2 = int(vec[4])
                    y2 = int(vec[5])
                    box = [x1, y1, x2, y2, conf, fname]

                    det_boxes[class_name].append(box)
        return det_boxes


class MetricsCOCO(object):
    def __init__(selfmin_iou=0.5, max_iou=0.95, iou_step=0.05):
        self.iou_thresholds = []
        iou = min_iou
        while iou <= max_iou:
            self.iou_thresholds.append(iou)
            iou += iou_step
    
    def mAP(self, gt_folder, det_folder):
        aps = []
        for thresh in self.iou_thresholds:
            metric = MetricsVOC(thresh)
            ap = metric.voc_ap(gt_folder, det_folder)
            aps.append(ap)
        return sum(aps) / len(aps)
